package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"

	"go-file-perception-model/internal/config"
	"go-file-perception-model/internal/logger"
)

type MCPRequest struct {
	JSONRPC string         `json:"jsonrpc"`
	ID      any            `json:"id"`
	Method  string         `json:"method"`
	Params  map[string]any `json:"params"`
}

type MCPResponse struct {
	JSONRPC string `json:"jsonrpc"`
	ID      any    `json:"id"`
	Result  any    `json:"result,omitempty"`
	Error   *struct {
		Code    int    `json:"code"`
		Message string `json:"message"`
	} `json:"error,omitempty"`
}

type Tool struct {
	Name        string         `json:"name"`
	Description string         `json:"description"`
	InputSchema map[string]any `json:"inputSchema"`
}

var mcpAPIKey string

func main() {
	// 初始化日志系统
	logger.InitLogger()
	
	// 加载配置
	cfg, err := config.LoadConfig()
	if err != nil {
		log.Fatalf("Failed to load config: %v", err)
	}
	
	// 添加调试信息
	log.Printf("Config loaded: %+v", cfg)
	log.Printf("MCP config: %+v", cfg.MCP)
	
	mcpAPIKey = cfg.MCP.APIKey
	log.Printf("MCP API key: '%s'", mcpAPIKey)
	log.Printf("MCP API key length: %d", len(mcpAPIKey))
	
	// 从配置中获取端口，如果未配置则使用默认端口8081
	port := cfg.MCP.Port
	if port == 0 {
		port = 8081
	}
	log.Printf("MCP port: %d", port)
	
	mux := http.NewServeMux()
	mux.HandleFunc("/mcp", authMiddleware(handleMCP))
	addr := fmt.Sprintf(":%d", port)
	log.Printf("Minimal Streamable HTTP MCP server listening on %s", addr)
	if err := http.ListenAndServe(addr, mux); err != nil {
		log.Fatal(err)
	}
}

// authMiddleware MCP请求鉴权中间件
func authMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// 从请求头中获取API密钥
		apiKey := r.Header.Get("X-API-Key")
		if apiKey == "" {
			// 如果请求头中没有，尝试从查询参数中获取
			apiKey = r.URL.Query().Get("api_key")
		}
		
		// 验证API密钥
		if apiKey != mcpAPIKey {
			w.Header().Set("Content-Type", "application/json; charset=utf-8")
			w.WriteHeader(http.StatusUnauthorized)
			_ = json.NewEncoder(w).Encode(MCPResponse{
				JSONRPC: "2.0",
				ID:      nil,
				Error: &struct {
					Code    int    `json:"code"`
					Message string `json:"message"`
				}{
					Code:    -32001,
					Message: "Unauthorized: Invalid API key",
				},
			})
			return
		}
		
		next(w, r)
	}
}

func handleMCP(w http.ResponseWriter, r *http.Request) {
	// 设置响应头，明确指定字符编码
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	
	defer r.Body.Close()
	var req MCPRequest
	decoder := json.NewDecoder(r.Body)
	if err := decoder.Decode(&req); err != nil {
		writeError(w, nil, -32700, "Parse error")
		return
	}
	switch req.Method {
	case "initialize":
		writeOK(w, req.ID, map[string]any{
			"protocolVersion": "2024-11-05",
			"capabilities":   map[string]any{"tools": map[string]any{}},
			"serverInfo":     map[string]any{"name": "Minimal MCP", "version": "0.1.0"},
		})
	case "tools/list":
		tools := []Tool{
			ToolSearchFilenamesDescriptor(), // search_filenames
			ToolViewFileContentDescriptor(), // view_file_content
			ToolRegexSearchDescriptor(),     // regex_search
			ToolSemanticSearchDescriptor(),  // semantic_search
		}
		writeOK(w, req.ID, map[string]any{"tools": tools})
	case "tools/call":
		params := req.Params
		name, _ := params["name"].(string)
		args, _ := params["arguments"].(map[string]any)
		stream, _ := params["stream"].(bool)
		if stream {
			w.Header().Set("Content-Type", "application/json; charset=utf-8")
			w.Header().Set("Cache-Control", "no-cache")
			flusher, ok := w.(http.Flusher)
			if !ok {
				writeError(w, req.ID, -32000, "Streaming not supported")
				return
			}
			enc := json.NewEncoder(w)
			_ = enc.Encode(map[string]any{"event": "start", "data": map[string]any{"tool": name}})
			flusher.Flush()
			var err error
			switch name {
			case "search_filenames":
				err = StreamSearchFilenames(r.Context(), args, enc, flusher)
			case "view_file_content":
				err = StreamViewFileContent(r.Context(), args, enc, flusher)
			case "regex_search":
				err = StreamRegexSearch(r.Context(), args, enc, flusher)
			case "semantic_search":
				err = StreamSemanticSearch(r.Context(), args, enc, flusher)
			default:
				err = fmt.Errorf("unknown tool: %s", name)
			}
			if err != nil {
				_ = enc.Encode(map[string]any{"event": "error", "data": map[string]any{"message": err.Error()}})
				flusher.Flush()
				return
			}
			_ = enc.Encode(map[string]any{"event": "end", "data": map[string]any{"ok": true}})
			flusher.Flush()
			return
		}
		var (
			res any
			err error
		)
		switch name {
		case "search_filenames":
			res, err = CallSearchFilenames(args)
		case "view_file_content":
			res, err = CallViewFileContent(args)
		case "regex_search":
			res, err = CallRegexSearch(args)
		case "semantic_search":
			res, err = CallSemanticSearch(args)
		default:
			err = fmt.Errorf("unknown tool: %s", name)
		}
		if err != nil {
			writeError(w, req.ID, -32603, err.Error())
			return
		}
		writeOK(w, req.ID, res)
	default:
		writeError(w, req.ID, -32601, "Method not found")
	}
}

func writeOK(w http.ResponseWriter, id any, result any) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	_ = json.NewEncoder(w).Encode(MCPResponse{JSONRPC: "2.0", ID: id, Result: result})
}

func writeError(w http.ResponseWriter, id any, code int, msg string) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	_ = json.NewEncoder(w).Encode(MCPResponse{JSONRPC: "2.0", ID: id, Error: &struct {
		Code    int    `json:"code"`
		Message string `json:"message"`
	}{Code: code, Message: msg}})
}

func callTool(ctx context.Context, name string, args map[string]any) (any, error) {
	switch name {
	case "search_filenames":
		return CallSearchFilenames(args)
	default:
		return nil, fmt.Errorf("unknown tool: %s", name)
	}
}

func streamCallTool(ctx context.Context, name string, args map[string]any, enc *json.Encoder, flusher http.Flusher) error {
	switch name {
	case "search_filenames":
		return StreamSearchFilenames(ctx, args, enc, flusher)
	case "view_file_content":
		return StreamViewFileContent(ctx, args, enc, flusher)
	case "regex_search":
		return StreamRegexSearch(ctx, args, enc, flusher)
	default:
		return fmt.Errorf("unknown tool: %s", name)
	}
}
